import random
import copy
from torch.utils.data import Dataset


class YoloDataset(Dataset):
    def __init__(self, dataset_dicts, map_func=None, train=True):
        super(YoloDataset, self).__init__()
        self.dataset_dicts = dataset_dicts
        self.train = train
        self.cut_min_offset = 0.2
        self.map_func = map_func

    def __len__(self):
        return len(self.dataset_dicts)

    def __getitem__(self, index):
        return self._get_train_item(index) if self.train else self._get_val_item(index)

    def _get_train_item(self, index):
        use_mosaic = random.randint(0, 1)
        if not use_mosaic:
            return self._get_val_item(index)
        else:
            image_indice = [index] + [random.randint(0, len(self.dataset_dicts) - 1) for i in range(3)]
            dataset_dict = None
            h, w = 0, 0
            cut_x, cut_y = 0, 0
            out_image = None
            for i, image_index in enumerate(image_indice):
                mapped_dataset_dict = self._get_val_item(image_index)
                if i == 0:
                    dataset_dict = copy.deepcopy(mapped_dataset_dict)
                    dataset_dict['annotations'] = []
                    h, w = dataset_dict['image'].shape[:2]
                    cut_x = random.randint(int(w * self.cut_min_offset), int(w * (1 - self.cut_min_offset)))
                    cut_y = random.randint(int(h * self.cut_min_offset), int(h * (1 - self.cut_min_offset)))
                    continue

                if i == 0:
                    cut_left, cut_top, cut_right, cut_bottom = 0, 0, cut_x, cut_y
                elif i == 1:
                    cut_left, cut_top, cut_right, cut_bottom = cut_x, 0, w, cut_y
                elif i == 2:
                    cut_left, cut_top, cut_right, cut_bottom = 0, cut_y, cut_x, h
                else:
                    cut_left, cut_top, cut_right, cut_bottom = cut_x, cut_y, w, h
                
                dataset_dict['image'][:, cut_top:cut_bottom, cut_left:cut_right] = mapped_dataset_dict['image'][:, cut_top:cut_bottom, cut_left:cut_right]

                for j in range(len(mapped_dataset_dict['annotations'])):
                    x1, y1, x2, y2 = mapped_dataset_dict['annotations'][j]['bbox']
                    x1 = max(min(x1, cut_right), cut_left)
                    x2 = max(min(x2, cut_right), cut_left)
                    y1 = max(min(y1, cut_bottom), cut_top)
                    y2 = max(min(y2, cut_bottom), cut_top)
                    if x1 < x2 and y1 < y2:
                        mapped_dataset_dict['annotations'][j]['bbox'] = [x1, y1, x2, y2]
                        dataset_dict['annotations'].append(mapped_dataset_dict['annotations'][j])

            return dataset_dict

    def _get_val_item(self, index):
        return self.map_func(self.dataset_dicts[index])
